# encoding: utf-8
import os
import h5py
import numpy as np
from keras.applications import imagenet_utils
from keras.preprocessing.image import ImageDataGenerator
from keras.applications import VGG16
from keras.layers.core import Lambda
from keras.layers import Input, GlobalAveragePooling2D, GlobalMaxPooling2D, Dense, Dropout
from keras.models import Model
from keras.callbacks import ModelCheckpoint, EarlyStopping
from keras.optimizers import adam, Adam
from sklearn.utils import shuffle
from keras.utils import to_categorical
from keras.utils.vis_utils import plot_model
from loss_history import LossHistory
preprocess = imagenet_utils.preprocess_input
from keras.models import load_model
import time
from keras import backend as K
K.clear_session()


# -------------------------------------
# 定义参数
# -------------------------------------
BATCH_SIZE = 4  # batch size
IMAGE_SIZE = (512, 512)  # 图片尺寸


# -------------------------------------
# 文件路径
# -------------------------------------
TEST_DATA_PATH = './data/valid/'
VGG16_MODEL = './model/vgg16-finetune-model.hdf5'
RESNET50_MODEL = './model/resnet50-finetune-model.hdf5'
INCEPTIONV3_MODEL = './model/inceptionv3-finetune-model.hdf5'
XCEPTION_MODEL = './model/xecption-finetune-model.hdf5'
MY_MODEL = './model/my-model.hdf5'
MODELS = [VGG16_MODEL, RESNET50_MODEL, INCEPTIONV3_MODEL, XCEPTION_MODEL, MY_MODEL]


def predict():

    BATCH_SIZE = 4
    valid_gen = ImageDataGenerator()
    test_generator = valid_gen.flow_from_directory(TEST_DATA_PATH, IMAGE_SIZE, shuffle=False, batch_size=BATCH_SIZE, color_mode='grayscale')
    labels = (test_generator.class_indices)
    label = dict((v,k) for k,v in labels.items())

    for MODEL in MODELS:

        model = load_model(MODEL)

        pred = model.predict_generator(test_generator, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=1)
        predicted_class_indices = np.argmax(pred, axis=1)
        predictions = [label[i] for i in predicted_class_indices]

        filenames = test_generator.filenames

        mm = 0
        mn = 0
        nm = 0
        nn = 0

        for idx, filename in enumerate(filenames):
            if filename[0] == 'm' and predictions[idx][0] == 'm':
                mm += 1
            elif filename[0] == 'm' and predictions[idx][0] == 'n':
                mn += 1
            elif filename[0] == 'n' and predictions[idx][0] == 'n':
                nn += 1
            elif filename[0] == 'n' and predictions[idx][0] == 'm':
                nm += 1

        total = mm + mn + nm + nn
        m_total = mm + mn
        n_total = nm + nn

        start = time.time()
        loss, _ = model.evaluate_generator(test_generator, max_queue_size=10, workers=1, use_multiprocessing=False, verbose=1)
        end = time.time()

        print('正确率：', (mm + nn) / total * 100)
        print('误报率：', nm / n_total * 100)
        print('漏报率:', mn / m_total * 100)
        print('Loss:', loss)
        print('时间开销(秒):', end - start)

        K.clear_session()


if __name__ == '__main__':
    predict()
